package top.ichengyi.springai.starter.factory;

import top.ichengyi.springai.starter.properties.MultiModelProperties;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.deepseek.DeepSeekChatModel;
import org.springframework.ai.deepseek.DeepSeekChatOptions;
import org.springframework.ai.deepseek.api.DeepSeekApi;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.zhipuai.ZhiPuAiChatModel;
import org.springframework.ai.zhipuai.ZhiPuAiChatOptions;
import org.springframework.ai.zhipuai.api.ZhiPuAiApi;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.stereotype.Component;

import java.util.HashMap;
import java.util.Map;

/**
 * AI模型工厂类
 */
@Component
@EnableConfigurationProperties(MultiModelProperties.class)
public class AiModelFactory {
    private final MultiModelProperties multiModelProperties;
    private final Map<String, ChatClient> clientCache = new HashMap<>();

    public AiModelFactory(MultiModelProperties multiModelProperties) {
        this.multiModelProperties = multiModelProperties;
    }

    public Map<String, ChatClient> buildAllClients() {
        Map<String, ChatClient> clients = new HashMap<>();
        for (MultiModelProperties.ModelConfig model : multiModelProperties.getModels()) {
            ChatClient client = buildClientFromModel(model);
            if (client != null) {
                clients.put(model.getName(), client);
            }
        }
        return clients;
    }

    private ChatClient buildClientFromModel(MultiModelProperties.ModelConfig modelConfig) {
        // 使用modelConfig的唯一标识符作为缓存键
        String cacheKey = modelConfig.getName() + ":" + modelConfig.getProvider();

        // 检查缓存中是否已存在该客户端实例
        if (clientCache.containsKey(cacheKey)) {
            return clientCache.get(cacheKey);
        }

        ChatClient client = null;
        // 根据提供商类型创建对应的客户端
        switch (modelConfig.getProvider().toLowerCase()) {
            case "openai":
                client = buildOpenAiClient(modelConfig);
                break;
            case "zhipu":
                client = buildZhipuClient(modelConfig);
                break;
            case "deepseek":
                client = buildDeepseekClient(modelConfig);
                break;
            default:
                System.out.println("不支持的提供商: " + modelConfig.getProvider());
        }

        // 将新创建的客户端放入缓存
        if (client != null) {
            clientCache.put(cacheKey, client);
        }

        return client;
    }


    private ChatClient buildDeepseekClient(MultiModelProperties.ModelConfig modelConfig) {
        DeepSeekChatOptions options = DeepSeekChatOptions.builder().model(modelConfig.getModel()).temperature(modelConfig.getTemperature()).build();
        DeepSeekApi api = DeepSeekApi.builder().apiKey(modelConfig.getApiKey()).build();
        DeepSeekChatModel chatModel = DeepSeekChatModel.builder().defaultOptions(options).deepSeekApi(api).build();
        return ChatClient.create(chatModel);
    }

    private ChatClient buildZhipuClient(MultiModelProperties.ModelConfig modelConfig) {
        ZhiPuAiChatOptions options = ZhiPuAiChatOptions.builder()
                .model(modelConfig.getModel())
                .temperature(modelConfig.getTemperature())
                .build();

        ZhiPuAiApi zhiPuAiApi = new ZhiPuAiApi(modelConfig.getApiKey());

        ZhiPuAiChatModel zhiPuAiChatModel = new ZhiPuAiChatModel(zhiPuAiApi, options);
        return ChatClient.create(zhiPuAiChatModel);
    }

    private ChatClient buildOpenAiClient(MultiModelProperties.ModelConfig modelConfig) {
        OpenAiChatOptions options = OpenAiChatOptions.builder()
                .model(modelConfig.getModel())
                .temperature(modelConfig.getTemperature())
                .build();

        OpenAiApi openAiApi = OpenAiApi.builder().apiKey(modelConfig.getApiKey()).baseUrl(modelConfig.getBaseUrl()).build();

        OpenAiChatModel chatModel = OpenAiChatModel.builder().defaultOptions(options).openAiApi(openAiApi).build();
        return ChatClient.create(chatModel);
    }

    public ChatClient getClientByName(String name) {
        Map<String, ChatClient> stringChatClientMap = buildAllClients();
        return stringChatClientMap.get(name);
    }
}